import copy
import datetime

from flask import current_app as app

from . import BaseHandler
from .utils import get_warning_id
from ...db.models.assets import Asset
from ...db.models.tasks import Task


class AssetHandler(BaseHandler):
    def handle(self):
        return self._handle_asset(response=self.response)

    @property
    def asset(self):
        asset_id = self.task.refId
        try:
            asset_obj = Asset.objects.with_id(asset_id)
        except Exception as e:
            app.logger.exception(f'[02]监测结果处理：asset {asset_id}Mongo记录查询失败:{e}')
        else:
            return asset_obj

    @property
    def asset_task(self):
        if not (result_id := self.asset.resultId):
            result_id = self.task.id
            self.asset.update(resultId=result_id)
            self.asset.reload()
        try:
            asset_task_obj = Task.objects.with_id(result_id)
        except Exception as e:
            app.logger.exception(f'[02]监测结果处理：asset-task {self.asset.resultId}关联Mongo记录查询失败:{e}')
        else:
            return asset_task_obj

    def _handle_asset(self, response):
        if not all([self.asset, self.asset_task]):
            return
        domain = self.task.domain
        ip2domain = response.get("ip2domain", {})
        subdomain_with_host = [sub for sub in response.get("subdomain", []) if sub.get("host", "").endswith('.' + domain)]
        ip_port = copy.deepcopy(response.get("port", {}))
        ret = {
            "ipBindDomains": [{"ip": ip2domain.get("ip"), "domains": ip2domain.get("data")}] if len(ip2domain.get("data", [])) else [],
            "subDomains": [
                {
                    "host": subdomain["host"],
                    "cname": subdomain["cname"],
                    "point": subdomain.get("ip") if subdomain.get("ip") != "-" else ""
                } for subdomain in subdomain_with_host],
            "ipPorts": [{"ip": ip_port.get("ip"), "ports": ip_port.get("data", [])}],
        }
        self._update_asset(ret)
        warnings = self.warning_check(asset_task=self.asset_task, ret=ret)
        ip_bind_domain = []
        for item in ret["ipBindDomains"]:
            domains = [d for d in item.get("domains", []) if d]
            ip_bind_domain.extend(domains)
        #     waf_list、alexa、basic_info、icp、ipdb、subdomain、http_get、http_method
        ret.update({
            "warnings": warnings,
            "stats": {
                "subdomainCount": len(ret["subDomains"]),
                "ipCount": len(ret["ipPorts"]),
                "portCount": sum((len(item.get("ports", [])) for item in ret["ipPorts"])),
                "ipBindDomainCount": sum(set(ip_bind_domain)),
                "addSubdomainCount": 0,
                "addIpCount": 0,
                "addPortCount": 0,
                "addIpBindDomainCount": 0,
                "removeSubdomainCount": 0,
                "removeIpCount": 0,
                "removePortCount": 0,
                "removeIpBindDomainCount": 0,
                "changePortServiceCount": 0,
            },
            "addition": {
                "target_ip": response.get("target_ip", ""),
                "waf_list": response.get("waf_list", []),
                "alexa": response.get("alexa", ''),
                "basic_info": response.get("basic_info", dict()),
                "icp": response.get("icp", dict()),
                "ipdb": response.get("ipdb", dict()),
                "subdomain": response.get("subdomain", dict()),
                "http_get": response.get("http_get", dict()),
                "http_method": response.get("http_method", ""),
                "port": response.get("port", {}),
            }
        })
        for warning in warnings:
            category = warning["category"]
            if category == "portService":
                key = "changePortServiceCount"
            else:
                op = warning["detail"]["op"]
                key = f"{op}{category.capitalize()}Count"
            ret["stats"][key] += 1
        return ret

    def _update_asset(self, update_set):
        try:
            Asset._get_collection().update_one(
                {
                    "_id": self.asset.id
                },
                {
                    "$set": {
                        **update_set,
                        "updateTime": datetime.datetime.now()}
                }
            )
        except Exception as e:
            app.logger.exception(f'[02]资产监测Worker更新资产asset {self.asset.id}失败:{e}')

    def warning_check(self, asset_task: Task, ret: dict):
        prev_result = asset_task.result.to_mongo()
        if not prev_result:
            return []

        warnings = []
        prev_subdomains_host = [item.get("host") for item in prev_result.get("subDomains", [])]
        next_subdomains_host = [item.get("host") for item in ret["subDomains"]]
        subdomain_changes = self.changes_compare(prev=prev_subdomains_host, new=next_subdomains_host)
        for domain in subdomain_changes["add"]:
            warning = self.format_warning(value=domain, op="add", category="subdomain")
            warnings.append(warning)

        for domain in subdomain_changes["remove"]:
            warning = self.format_warning(value=domain, op="remove", category="subdomain")
            warnings.append(warning)

        prev_ip_ports = prev_result.get("ipPorts", [{}])
        prev_ip_ports = prev_ip_ports[0].get("ports", {}) if prev_ip_ports else {}
        curr_ip_ports = ret["ipPorts"][0].get("ports", {})

        prev_port = [item.get("port") for item in prev_ip_ports]
        next_port = [item.get("port") for item in curr_ip_ports]
        port_changes = self.changes_compare(prev=prev_port, new=next_port)
        for port in port_changes["add"]:
            warning = self.format_warning(value=port, op="add", category="port")
            warnings.append(warning)

        for port in port_changes["remove"]:
            warning = self.format_warning(value=port, op="remove", category="port")
            warnings.append(warning)

        if len(port_changes["stay"]):
            current_ports = ret.get("ipPorts", [{}])[0].get("ports", [])
            current_ports_service_map = {}
            for port in current_ports:
                current_ports_service_map[port.get("port")] = port.get("service") or "未知服务"

            prev_port_service_map = {}
            prev_ports = prev_result.get("ipPorts", [{}])[0].get("ports", [])
            for port in prev_ports:
                prev_port_service_map[port.get("port")] = port.get("service") or "未知服务"

            for port in port_changes["stay"]:
                if prev_port_service_map[port] == current_ports_service_map[port]:
                    continue
                warning = {
                    "id": get_warning_id(
                        task_type="asset",
                        job_id=self.task.jobId,
                        url=port,
                        risk_title="portService"
                    ),
                    "level": 2,
                    "title": f"端口服务{port}({prev_port_service_map[port]})调整为{port}({current_ports_service_map[port]})",
                    "category": 'portService',
                    "affects": port,
                    "detail": {
                        "prev": prev_port_service_map[port],
                        "current": current_ports_service_map[port],
                    }
                }
                warnings.append(warning)
        return warnings

    def changes_compare(self, prev, new):
        remove_map = {key: True for key in prev}
        next_map = {key: True for key in new}
        add_map = {}
        stay_map = {}
        for key in next_map:
            if key not in remove_map:
                add_map[key] = True
            else:
                stay_map[key] = True
                del remove_map[key]
        return {
            "add": list(add_map.keys()),
            "remove": list(remove_map.keys()),
            "stay": list(stay_map.keys())
        }

    def format_warning(self, value, op, category):
        op_mapping = {
            "subdomain": {
                "add": "上线",
                "remove": "下线"
            },
            "port": {
                "add": "开放",
                "remove": "关闭"
            }
        }
        category_title_mapping = {
            "subdomain": f"[{value}]子域名{op_mapping[category][op]}",
            "port": f"[{value}]端口{op_mapping[category][op]}`"
        }
        warning = {
            "id": get_warning_id(
                task_type="asset",
                job_id=self.task.jobId,
                url=value,
                risk_title=category
            ),
            "level": 2,
            "title": category_title_mapping[category],
            "category": category,
            "affects": value,
            "detail": {
                "op": op
            },
        }
        return warning
